use core::{ops::Deref, ptr, u32};

use stm32f4::stm32f407::{self, flash};

pub struct FLASH;

impl Deref for FLASH {
    type Target = flash::RegisterBlock;
    fn deref(&self) -> &Self::Target {
        unsafe { &*stm32f407::FLASH::PTR }
    }
}

pub const SECTOR0_ADDR: u32 = 0x08000000;
pub const SECTOR0_SIZE: u32 = 16 * 1024;

pub const SECTOR1_ADDR: u32 = SECTOR0_ADDR + SECTOR0_SIZE;
pub const SECTOR1_SIZE: u32 = 16 * 1024;

pub const SECTOR2_ADDR: u32 = SECTOR1_ADDR + SECTOR1_SIZE;
pub const SECTOR2_SIZE: u32 = 16 * 1024;

pub const SECTOR3_ADDR: u32 = SECTOR2_ADDR + SECTOR2_SIZE;
pub const SECTOR3_SIZE: u32 = 16 * 1024;

pub const SECTOR4_ADDR: u32 = SECTOR3_ADDR + SECTOR3_SIZE;
pub const SECTOR4_SIZE: u32 = 64 * 1024;

pub const SECTOR5_ADDR: u32 = SECTOR4_ADDR + SECTOR4_SIZE;
pub const SECTOR5_SIZE: u32 = 128 * 1024;

pub const SECTOR6_ADDR: u32 = SECTOR5_ADDR + SECTOR5_SIZE;
pub const SECTOR6_SIZE: u32 = 128 * 1024;

pub const SECTOR7_ADDR: u32 = SECTOR6_ADDR + SECTOR6_SIZE;
pub const SECTOR7_SIZE: u32 = 128 * 1024;

#[derive(Debug)]
pub enum Error {
    INDEX,
    PGS,
    PGP,
    PGA,
    WRP,
    WRITE,
    ID,
}

struct Sector {
    addr: u32,
    size: u32,
}

impl Sector {
    pub const fn new(addr: u32, size: u32) -> Self {
        Self { addr, size }
    }
}

const SECTORS: [Sector; 8] = [
    Sector::new(SECTOR0_ADDR, SECTOR0_SIZE),
    Sector::new(SECTOR1_ADDR, SECTOR1_SIZE),
    Sector::new(SECTOR2_ADDR, SECTOR2_SIZE),
    Sector::new(SECTOR3_ADDR, SECTOR3_SIZE),
    Sector::new(SECTOR4_ADDR, SECTOR4_SIZE),
    Sector::new(SECTOR5_ADDR, SECTOR5_SIZE),
    Sector::new(SECTOR6_ADDR, SECTOR6_SIZE),
    Sector::new(SECTOR7_ADDR, SECTOR7_SIZE),
];

fn find_sector_id(addr: u32) -> Result<u8, Error> {
    for (index, sector) in SECTORS.iter().enumerate() {
        if (addr >= sector.addr) && (addr < (sector.addr + sector.size)) {
            return Ok(index as u8);
        }
    }
    Err(Error::INDEX)
}

struct LockGuard {}

impl Drop for LockGuard {
    fn drop(&mut self) {
        FLASH.cr.modify(|_, w| w.lock().locked());
    }
}

fn unlock() -> LockGuard {
    FLASH.keyr.write(|w| w.key().variant(0x45670123));
    FLASH.keyr.write(|w| w.key().variant(0xCDEF89AB));
    LockGuard {}
}

struct DCGuard {}

impl Drop for DCGuard {
    fn drop(&mut self) {
        FLASH.acr.modify(|_, w| w.dcen().enabled());
    }
}

fn disable_dc() -> DCGuard {
    FLASH.acr.modify(|_, w| w.dcen().disabled());
    DCGuard {  }
}

struct SERGuard {}

impl Drop for SERGuard {
    fn drop(&mut self) {
        FLASH.cr.modify(|_, w| w.ser().clear_bit());
    }
}

fn start_ser() -> SERGuard {
    FLASH.cr.modify(|_, w|w.ser().sector_erase());
    SERGuard {  }
}

struct PGGuard {}

impl Drop for PGGuard {
    fn drop(&mut self) {
        FLASH.cr.modify(|_, w| w.pg().clear_bit());
    }
}

fn start_pg() -> PGGuard {
    FLASH.cr.modify(|_, w| w.pg().program());
    PGGuard {  }
}

fn wait_done() -> Result<(), Error> {
    loop {
        let sr = FLASH.sr.read();
        if sr.pgserr().bit_is_set() {
            return Err(Error::PGS);
        }
        if sr.pgperr().bit_is_set() {
            return Err(Error::PGP);
        }
        if sr.pgaerr().bit_is_set() {
            return Err(Error::PGA);
        }
        if sr.wrperr().bit_is_set() {
            return Err(Error::WRP);
        }
        if sr.bsy().bit_is_clear() {
            break;
        }
    }

    Ok(())
}

fn sector_erase_with_id(id: u8) -> Result<(), Error> {
    let _dc = disable_dc();
    wait_done()?;
    FLASH.cr.modify(|_, w| w.psize().psize32().snb().variant(id));
    let _ser = start_ser();
    FLASH.cr.modify(|_, w| w.strt().start());
    wait_done()?;
    Ok(())
}

fn write_u32(addr: u32, value: u32) -> Result<(), Error> {

    wait_done()?;
    FLASH.cr.modify(|_, w| w.psize().psize32());
    
    unsafe {
        let _pg = start_pg();
        (addr as *mut u32).write_volatile(value);
        wait_done()?;
    }
    unsafe {
        if (addr as *const u32).read_volatile() != value {
            return Err(Error::WRITE);
        }
    }

    Ok(())
}

pub fn erase(addr: u32, size: u32) -> Result<(), Error> {
    let _lock = unlock();

    let mut index = 0;
    let mut id = find_sector_id(addr)?;
    
    while index < size {
        if (id as usize) >= SECTORS.len() {
            return Err(Error::ID);
        }
        sector_erase_with_id(id)?;
        index += SECTORS[id as usize].size;
        id += 1;
    }

    Ok(())
}

pub fn write(addr: u32, buf: &[u8]) -> Result<(), Error> {
    let _lock = unlock();
    let mut index = 0;
    let ptr = buf.as_ptr();
    while index < buf.len() {
        
        let value: u32 = unsafe {
            u32::from_ne_bytes([
                ptr.add(index).read(),
                ptr.add(index + 1).read(),
                ptr.add(index + 2).read(),
                ptr.add(index + 3).read()
            ])
        };
        write_u32(addr + (index as u32), value)?;
        index += 4;
    }
    Ok(())
}

pub fn read(addr: u32, buf: &mut[u8]) {
    unsafe {
        ptr::copy_nonoverlapping(
            addr as *const u8, 
            buf.as_mut_ptr(), 
            buf.len()
        );
    }
}

